# coding: utf-8
# Author：WangTianRui
# Date ：2022/3/5 15:53
import sys, os, json

sys.path.append(os.path.abspath(os.path.join(os.getcwd(), "../")))
# import Model_voicefilter.SR_16K.model as model
import torch
import utils.features as features_utils
import utils.utils_ as utils
import copy, pickle
import numpy as np
from tqdm import tqdm
import soundfile as sf


def get_wavs(root):
    all_wavs = utils.get_all_wavs(root)
    all_infos = {}
    for i, item in enumerate(all_wavs):
        if item.split("\\")[-2] not in all_infos.keys():
            all_infos[item.split("\\")[-2]] = [item]
        else:
            all_infos[item.split("\\")[-2]].append(item)
    return all_infos


def get_embedding(model, wav_path, save_path):
    sig, sr = sf.read(wav_path)
    if sr != 16000:
        sig = features_utils.librosa.resample(sig, orig_sr=sr, target_sr=16000).astype(np.float32)
    embedding = model(torch.tensor([sig]))[0].detach().numpy()
    print(save_path + ";" + str(embedding.shape))
    np.save(save_path, embedding)


if __name__ == '__main__':
    import Model_ResNetSE34.ResNetSE34V2 as ResNetSE34V2

    root = r"/root/wangtianrui/dataset/jiutian/dns_2022/datasets_fullband/clean_fullband/read_speech"
    save_root = r"/root/wangtianrui/dataset/jiutian/dns_2022/pdns/read_speech_resnetse34v2"
    embedder_pt = torch.load(
        r"/root/wangtianrui/dataset/jiutian/dns_2022/pdns/save_logs/voicefilterSR/SR_models/baseline_v2_ap.model",
        map_location="cpu")
    embedder = ResNetSE34V2.ResNetSE()
    for key in list(embedder_pt.keys()):
        if str(key).startswith('__S__'):
            embedder_pt[key.replace("__S__.", "")] = embedder_pt[key]
            embedder_pt.pop(key, '404')
        if str(key).startswith('__L__'):
            embedder_pt.pop(key, '404')
    embedder.load_state_dict(embedder_pt)
    embedder.eval()

    all_json = json.load(open(r"/root/wangtianrui/dataset/jiutian/dns_2022/pdns/48k/speaker_json.json", "r"))
    for key in all_json.keys():
        all_wavs = all_json[key]
        for item in all_wavs:
            npy_path = os.path.join(save_root, item + ".npy")
            if not os.path.exists(npy_path):
                wav_path = os.path.join(root, item)
                print(wav_path, npy_path)
                get_embedding(embedder, wav_path, npy_path)
